#!/usr/bin/python
#coding:utf-8

# Copyright 2011 Nicolau Leal Werneck, Anna Helena Reali Costa and
# Universidade de São Paulo
#
# Licensed under theglim Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

###############################################################################
## Open a (pinhole camera model) picture, find its orientation and
## extract the rectified.
##
## Changes that must be performed as soon as possible: read the
## intrinsic parameters from somewhere, and make it easy to switch to
## the equirectangular model.

import sys

import os


def utime():
    '''Pick user time from os to measure the code performance.'''
    return os.times()[0]



if __name__ == '__main__':
    
    import matplotlib
    if sys.argv[0][-7:] == '-nop.py':
        PlotStuff=False
    else:
        PlotStuff=True
        from pylab import *
        import plot_aux
        if sys.argv[0][-7:] == '-nox.py':
            matplotlib.use('Agg') 


from numpy import *

import scipy.io

from numpy import dot

# from camori import PicturePinhole, PictureEqrec, quaternion_to_matrix, measure_error, quaternion_product, dir_colors
from camori_equidistant import PictureEquidistant
from camori_aux import *
from quaternion import Quat, random_quaternion

import simplejson

import scipy.optimize
#from fmin_sa import Sampling
from simplex_sa import SimplexSO3

import Image

set_printoptions(precision=7)



import filterSQP

######################################################################
## Functions used in the optimization

## Conic constraint
def val_c(x):
    return np.linalg.norm(x)
def grad_c(x):
    return x/np.linalg.norm(x)
def hess_c(x):
    nx = np.linalg.norm(x)
    return (nx**2 * np.identity(x.shape[0]) - np.outer(x,x)) / nx**3

## Target function
def val_f(x,*fargs):
    return camori_angle_error(x, fargs[0],fargs[1],fargs[2])
def grad_f(x,*fargs):
    return camori_angle_error_gradient(x, fargs[0],fargs[1],fargs[2])
def hess_f(x,*fargs):
    return camori_angle_error_hessian(x, fargs[0],fargs[1],fargs[2])
##
######################################################################

def create_pic(job_params, framenum):
    filename = job_params['root_directory']+'/frames/'+job_params['filename_format']%framenum

    im = Image.open(filename)
    frame = array(im.convert('RGB'), dtype=float)
    imr = array(im.convert('RGB'), dtype=float)
    imr = imr[:,:,:3] #remove alpha channel


    # Smooth out
    if ("gaussian_smoothing_factor" in job_params.keys() and 
        job_params["gaussian_smoothing_factor"] > 0):
        for c in range(3):
            imr[:,:,c] = scipy.ndimage.gaussian_filter( imr[:,:,c], \
                          double(job_params["gaussian_smoothing_factor"]))

    ## Creates picture object
    if not 'projection_model' in job_params:
        raise "Missing camra model in job file"
    else:
        model = job_params['projection_model']

    if not (model in ['pinhole', 'harris', 'polar_equidistant', 'cylindrical_equidistant']):
        raise NotImplementedError

    elif model == 'pinhole':
        ## Intrinsic parameters
        focal_distance = job_params['focal_distance']
        p_point = array(job_params['principal_point'])
        i_param = array([0.0, focal_distance, p_point[0], p_point[1]])

    elif model == 'harris':
        ## Intrinsic parameters
        focal_distance = job_params['focal_distance']
        p_point = array(job_params['principal_point'])
        distortion = job_params['distortion_coefficient']
        i_param = array([2.0, focal_distance, p_point[0], p_point[1], distortion])
        
    elif model == 'polar_equidistant':
        ## Intrinsic parameters
        focal_distance = job_params['focal_distance']
        p_point = array(job_params['principal_point'])
        i_param = array([3.0, focal_distance, p_point[0], p_point[1]])

    elif model == 'cylindrical_equidistant':
        ## Intrinsic parameters
        focal_distance = job_params['focal_distance']
        p_point = array(job_params['principal_point'])
        i_param = array([4.0, focal_distance, p_point[0], p_point[1]])

    sqp_funcs = (val_c, grad_c, hess_c, val_f, grad_f, hess_f)
    pic = PictureEquidistant(imr, i_param)


    pic.framenumber = framenumber


    return pic

if __name__ == '__main__':
    tt_total_init = utime()
    if PlotStuff:
        rc('text',usetex=False)

    ## Avoid zero divide warnins...
    np.seterr(divide='ignore')

    if PlotStuff:
        ## Plot stuff immediately
        ion()

    
    ##################################################################
    ## Read all the job parameters
    if len(sys.argv) != 2 and len(sys.argv) != 4:
        print sys.argv[0], '<job_file.json> [first frame last_frame]'
        raise Exception('Insufficient number of parameters')
        
    finput = open(sys.argv[1])
    job_params = simplejson.load(finput)
    finput.close()

    fileroot = job_params['root_directory']


    ## Edgel extractor parameters
    gmethod = job_params['edge_detection_method']
    gspc = job_params['grid_spacing']
    glim = job_params['gradient_threshold']

    do_robust = job_params['do_robust']
    initial_trials = job_params['initial_trials']

    do_optimization = job_params['do_optimization']
    optimization_tolerance = job_params['optimization_tolerance']

    if 'do_decimation' in job_params.keys():
        do_decimation = job_params['do_decimation']
        dec_t = job_params['decimator_threshold']
        dec_d = job_params['decimator_distance']
        dec_m = job_params['decimator_method']
        dec_l = job_params['decimator_lower']
    else:
        do_decimation = 0

    if 'do_multiple_initializations' in job_params.keys():
        do_multiple_initializations = job_params['do_multiple_initializations']
    else:
        do_multiple_initializations = 0

    if 'do_second_optimization' in job_params.keys():
        do_second_optimization = job_params['do_second_optimization']
    else:
        do_second_optimization = 0


    ##
    #################################################################


    #################################################################
    ## Load images and initialize pic objects
    tt_init = utime()

    ## Sets filename list from input argument
    if len(sys.argv) == 4:
        frmf = int(sys.argv[2])
        frml = int(sys.argv[3])
        the_frames = range(frmf, frml+1)
    if len(sys.argv) == 2:
        the_frames = [int(x) for x in sys.stdin.read().split()]

    picL = []
    for framenumber in the_frames:
        #sys.stderr.write('Reading frame %d\n'%framenumber)
        print 'Reading frame', framenumber
        pic = create_pic(job_params, framenumber)
        ## Extract the edgels from the image using the grid mask
        pic.extract_edgels(gspc, glim, method=gmethod)
        ## Caculate the edgel (interpretation plane) normals. Necessary
        ## for the random search (we might calculate the normals inside
        ## there...)
        pic.calculate_edgel_normals()
        ## Filter the edgels to remove the lonely ones. Impose a kind of
        ## curvature restriction in the edgels, but still not a full
        ## countour extraction...  Uses a kd-tree, but still slow.
        if do_decimation and dec_t > 0:
            pic.decimate_edgels(dec_t, dec_d, dec_l, dec_m)

        ## Jacobians 
        pic.calculate_edgel_jacobians()

        picL.append(pic)
    tt_image_initialization = utime()-tt_init
    ##
    ##################################################################

    ##################################################################
    ## Calculate initial estimates
    tt_init = utime()
    for pic in picL:
        fp = array(job_params["fp_robust"])
        pic.qini,bpk_a,bpk_b,bpk_c = pic.random_search(initial_trials, fp)
        pic.qini = pic.qini.canonical()
    tt_initial_estimate = utime() - tt_init
    ##
    ##################################################################

    #################################################################
    ## Perform optimization procedure for calibration, based on
    ## "local" optimizations to find orientation for each image and
    ## for each sampled combination of intrinsic parameters.
    tt_init = utime()

    def calibration_target(pp, picL):
        # sys.stderr.write('%f %f '%(pp[0],pp[1]))
        print 70*'='
        print pp
        errs = []
        for n,pic in enumerate(picL):
            fp = array(job_params["fp_optimization"])
            sqp_funcs = (val_c, grad_c, hess_c, val_f, grad_f, hess_f)
            i_param = copy(pic.i_param)
            i_param[1] = pp
            args_f = (pic.edgels, i_param, fp)
            filterSQPout = filterSQP.filterSQP(pic.qini.q,.0,1e-3,sqp_funcs,args_f,
                                               delta_tol=optimization_tolerance)
            xo, err, sqp_iters,Llam,Lrho = filterSQPout
            errs.append(err)

            qopt = Quat(xo)
            pic.lastq = qopt
            pic.lasterr = err
            # sys.stderr.write('%d'%n, pic.qini, qopt, err
            print n, pic.qini, qopt, err
        #sys.stderr.write(' %f\n'%tot_err)
        print errs
        return errs



    pprez = 0.01
    xx = mgrid[0.6:1.8+pprez:pprez]

    pp0 = array([pic.fd,])

    xx = xx * pic.fd

    Npp = xx.shape[0]

    Fpp = zeros((Npp,len(the_frames)))

    for npp in range(Npp):
        Fpp[npp] = array(calibration_target(xx[npp], picL))
        print xx[npp], Fpp[npp]

    tt_filtersqp = utime() - tt_init
    ##
    ##################################################################

    tt_total = utime() - tt_total_init
    print 'total time:', tt_total

    savez(sys.argv[1]+'-focal-one', dist_focal=xx, func_obj=Fpp)
    plot(xx, Fpp/Fpp.min())
    grid()

